import keyboard
import time
from env.airsim_gym3 import AirSimMultiAgentEnv
import numpy as np
from configs.config import Config

def visualize_distance_sensors(distances):
    """可视化距离传感器数据"""
    if not distances or all(d is None for d in distances):
        print("  无有效距离传感器数据")
        return

    # 16个传感器均匀分布在车辆周围
    sensor_positions = [f"{i * 22.5}°" for i in range(16)]  # 每个传感器间隔22.5度

    # 找到最近的物体及其方向
    min_distance = float('inf')
    min_idx = -1
    for i, d in enumerate(distances):
        if d is not None and d < min_distance:
            min_distance = d
            min_idx = i

    # 打印传感器数据的简单可视化
    print("  距离传感器可视化 (角度: 距离):")
    for i, (pos, d) in enumerate(zip(sensor_positions, distances)):
        if d is None:
            status = "N/A"
        elif d < 1.0:  # 近距离（红色警报）
            status = f"\033[91m{d:.2f}m\033[0m"  # 红色文本
        elif d < 3.0:  # 中距离（黄色警告）
            status = f"\033[93m{d:.2f}m\033[0m"  # 黄色文本
        else:  # 远距离（绿色安全）
            status = f"\033[92m{d:.2f}m\033[0m"  # 绿色文本

        # 标记最近的传感器
        marker = " <<< 最近" if i == min_idx else ""
        print(f"  {pos}: {status}{marker}")

    # 打印障碍物摘要
    if min_idx != -1:
        print(f"  最近障碍物: {min_distance:.2f}m, 方向: {sensor_positions[min_idx]}")
    else:
        print("  未检测到障碍物")


def test_keyboard_control():
    # 创建环境实例
    env = AirSimMultiAgentEnv()

    # 确保只使用一个智能体
    env.num_agents = 1

    try:
        # # 定义目标位置字典
        target_pos = env.config.get("target_pos", None)

        # target_pos = np.asarray([50, 0, -0.61])
        # 重置环境，设置目标坐标
        obs = env.reset()

        # print(f"目标坐标: {target_pos}")
        print("键盘控制说明:")
        print("  W: 加速")
        print("  S: 减速/倒车")
        print("  A: 左转")
        print("  D: 右转")
        print("  空格: 刹车")
        print("  R: 重置环境")
        print("  Q: 退出")

        # 初始化动作
        action = np.array([0.0, 0.0, 0.0])  # 油门、转向、刹车
        running = True

        while running:
            # 处理键盘输入
            if keyboard.is_pressed('w'):
                action[0] = 1.0  # 前进
            elif keyboard.is_pressed('s'):
                action[0] = -1.0  # 后退
            else:
                action[0] = 0.0  # 停止

            if keyboard.is_pressed('a'):
                action[1] = -1.0  # 左转
            elif keyboard.is_pressed('d'):
                action[1] = 1.0  # 右转
            else:
                action[1] = 0.0  # 回正

            if keyboard.is_pressed('space'):
                action[2] = 1.0  # 刹车
            else:
                action[2] = 0.0  # 松开刹车

            if keyboard.is_pressed('r'):
                print("\n重置环境...")
                obs = env.reset()
                time.sleep(1)  # 等待重置完成

            if keyboard.is_pressed('q'):
                running = False  # 退出程序

            # 执行动作
            next_obs, rewards, dones, infos = env.step([action])

            # 打印控制状态
            # print(f"\n控制状态: 油门={action[0]:.2f}, 转向={action[1]:.2f}, 刹车={action[2]:.2f}")

            # 打印智能体信息
            info = infos["infos"][0]
            # print(f"\n智能体信息:")
            print(f"奖励: {rewards[0]:.4f}")
            print(f"自身坐标: ({info['x']:.2f}, {info['y']:.2f}, {info['z']:.2f})")
            # print(f"速度: {info['speed']:.2f} m/s")
            print(f"距离目标点: {info['distance']:.2f} m")

            # 可视化距离传感器数据
            # print("\n距离传感器数据:")
            visualize_distance_sensors(info['distance_sensor_data'])

            # 检查是否完成任务
            if dones[0]:
                print("\n任务完成!")
                time.sleep(2)
                # 重置环境
                obs = env.reset()

            # 控制更新频率
            time.sleep(0.1)

    except Exception as e:
        print(f"测试过程中出现异常: {e}")
    finally:
        # 关闭环境
        env.close()


if __name__ == "__main__":
    test_keyboard_control()
